import asyncio
from elasticsearch._async.helpers import async_scan, async_bulk

from core.es import Index_KJJ_ALUMNI
from utils import strutil
from utils.es import es_client
from utils.iterutil import generate_batch


async def get_all_alumni():
    query = {"query": {"match_all": {}}}
    items = async_scan(es_client, query, index=Index_KJJ_ALUMNI)
    async for item in items:
        source = item['_source']
        source['id'] = item['_id']
        source['score'] = item['_score']
        yield source


async def bulk_update_docs(docs, stats_only=False, raise_on_error=True):
    bulk_actions = list()
    for doc in docs:
        bulk_actions.append({
            '_op_type': 'update',
            '_index': Index_KJJ_ALUMNI,
            '_id': doc.pop('id'),
            'doc': doc,
        })
    return await async_bulk(es_client, bulk_actions, stats_only=stats_only, raise_on_error=raise_on_error)


async def execute():
    docs = []
    async for alumni in get_all_alumni():
        name = alumni.get('name_zh', alumni.get('name'))
        if not name or not strutil.contain_zh(name):
            continue
        docs.append({"id": alumni["id"], "name_pinyin": strutil.cn_name2pinyin(name)})
    for batch in generate_batch(docs):
        result = await bulk_update_docs(batch, stats_only=True)
        print(result)

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    loop.run_until_complete(execute())
